# 导入需要的包
import sys
from matplotlib import pyplot
from keras.models import Sequential
from keras.layers import Conv2D
from keras.layers import MaxPooling2D
from keras.layers import Dense
from keras.layers import Flatten
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import tensorflow as tf
from keras.optimizers import SGD

# 进行你的 TensorFlow 操作
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)

print(tf.__version__)
device_name  = tf.test.gpu_device_name()
print("gpu设备名:",device_name)
print("设置可见设备",tf.config.experimental.set_visible_devices)
print('GPU:', tf.config.list_physical_devices('GPU'))
print("是否GPU可用",tf.test.is_gpu_available())
# 输出可用的GPU数量
print("可用的GPU数量:", len(tf.config.experimental.list_physical_devices('GPU')))
tf.device('GPU:1')

# 创建一个cnn模型
def define_cnn_model():

    # 使用序列模型
    model = Sequential()
    # 卷积层
    # 添加激活函数relu，增加非线性表达能力，如果没有激活函数，那么神经网络只能表示一种线性变换
    # 卷积核初始化方法为he_uniform
    # padding 使得卷积操作后的输出数据维度与输入数据维度相同。
    # input_shape=(200, 200, 3) 描述输入数据，高度x宽度（即矩阵的行数和列树），3为通道数
    # 结果：一次卷积后会的到一个（200，200，32）的特征体
    model.add(Conv2D(32, (3, 3), activation="relu",
                     kernel_initializer='he_uniform',
                     padding="same",
                     input_shape=(200, 200, 3),
                     ))
    # 最大池化层
    model.add(MaxPooling2D((2, 2)))

    model.add(Conv2D(64, (3, 3), activation="relu",
                     kernel_initializer='he_uniform',
                     padding="same",
                     ))
    model.add(MaxPooling2D((2, 2)))

    model.add(Conv2D(128, (3, 3), activation="relu",
                     kernel_initializer='he_uniform',
                     padding="same",
                     ))
    model.add(MaxPooling2D((2, 2)))
    # Flatten层
    model.add(Flatten())
    # 全连接层
    model.add(Dense(128, activation="relu", kernel_initializer='he_uniform'))
    model.add(Dense(1, activation="sigmoid"))  # 输出层0，1，sigmoid模型实现输出值0~1之间，分别代表猫狗

    
    # 优化器，随机梯度下降，为模型找到最佳的参数
    # 学习率和动量
    opt = SGD(lr=0.001, momentum=0.9)
    # 编译模型
    # 损失函数为二元交叉熵 性能指标为准确率
    model.compile(optimizer=opt,
                  loss='binary_crossentropy',
                  metrics=['accuracy'])
    return model


# 训练模型
def train_cnn_model():

    # 实例化模型
    model = define_cnn_model()
    # 创建图片生成器,产生图片并输入
    # ImageDataGenerator的flow_from_directory()方法可以自动识别子目录并将它们作为不同的类别，
    # 并且按照字母顺序给子目录分配类别标签
    datagen = ImageDataGenerator(rescale=1.0 / 225.0)
    train_it = datagen.flow_from_directory(
       'G:\\Research\\Test_data\\cat_dag_database\\dogs_cats\\data\\train',
        class_mode='binary',
        batch_size=64,  # 一次产生并输入64张图片
        target_size=(200, 200)  # 缩放图片为200x200，和输入图片大小相同
    )
    print("hello")
    # 训练模型
    model.fit(train_it,
            steps_per_epoch=len(train_it),
            epochs=10,
            verbose=1)

    # 把模型保存到文件夹
    model.save("G:\\ALL_Project\\Python_Project\\Data_clean\\my_model.h5")



train_cnn_model()